import os
import pandas as pd
from langchain.prompts import PromptTemplate
from langchain.schema import BaseOutputParser
from langchain.chains import LLMChain
from langchain_openai import ChatOpenAI
from typing import List, Dict
import json
import re
import glob
from dotenv import load_dotenv

load_dotenv()

# 阿里云百炼配置
DASHSCOPE_CONFIG = {
    "api_key": os.getenv('DASHSCOPE_API_KEY'),
    "base_url": os.getenv('BASE_URL'),
    "model": "qwen-long"
}

# 验证API密钥
if not DASHSCOPE_CONFIG["api_key"]:
    raise ValueError("请设置DASHSCOPE_API_KEY环境变量")


class StatisticalQAOutputParser(BaseOutputParser):
    """解析LLM输出的统计问答对"""

    def parse(self, text: str) -> List[Dict[str, str]]:
        """解析统计问答对文本"""
        qa_pairs = []

        # 预处理：提取JSON内容（处理代码块格式）
        json_text = self._extract_json_content(text)

        # 尝试解析JSON格式
        try:
            if json_text.strip().startswith('[') or json_text.strip().startswith('{'):
                data = json.loads(json_text)
                if isinstance(data, list):
                    return data
                elif isinstance(data, dict) and 'qa_pairs' in data:
                    return data['qa_pairs']
        except json.JSONDecodeError:
            pass

        # 解析文本格式的问答对
        lines = text.strip().split('\n')
        current_qa = {}

        for line in lines:
            line = line.strip()
            if not line:
                continue

            # 匹配问题格式
            q_match = re.match(r'^(?:问题?\d*|Q\d*|Question\d*)\s*[：:：]?\s*(.+)$', line, re.IGNORECASE)
            if q_match:
                if current_qa and 'question' in current_qa and 'answer' in current_qa:
                    qa_pairs.append(current_qa)
                clean_question = self._clean_content(q_match.group(1))
                current_qa = {'question': clean_question}
                continue

            # 匹配答案格式
            a_match = re.match(r'^(?:答案?\d*|A\d*|Answer\d*)\s*[：:：]?\s*(.+)$', line, re.IGNORECASE)
            if a_match and 'question' in current_qa:
                clean_answer = self._clean_content(a_match.group(1))
                current_qa['answer'] = clean_answer
                continue

            # 如果当前有问题但没有答案，将此行作为答案
            if 'question' in current_qa and 'answer' not in current_qa:
                clean_answer = self._clean_content(line)
                current_qa['answer'] = clean_answer

        # 添加最后一个问答对
        if current_qa and 'question' in current_qa and 'answer' in current_qa:
            qa_pairs.append(current_qa)

        return qa_pairs

    def _extract_json_content(self, text: str) -> str:
        """提取JSON内容，处理代码块格式"""
        text = text.strip()

        # 检查是否包含```json代码块
        json_block_pattern = r'```json\s*\n?([\s\S]*?)\n?```'
        json_match = re.search(json_block_pattern, text, re.IGNORECASE)

        if json_match:
            return json_match.group(1).strip()

        return text

    def _clean_content(self, content: str) -> str:
        """清理内容，移除标号和多余的冒号"""
        content = re.sub(r'^\d+[.、：:：]?\s*', '', content)
        content = re.sub(r'^[A-Za-z][.、：:：]?\s*', '', content)
        content = re.sub(r'^[：:：]+\s*', '', content)
        return content.strip()


class CSVStatisticalQAGenerator:
    """CSV统计问答对生成器"""

    def __init__(self):
        # 初始化阿里云百炼模型
        self.llm = ChatOpenAI(
            api_key=DASHSCOPE_CONFIG["api_key"],
            base_url=DASHSCOPE_CONFIG["base_url"],
            model=DASHSCOPE_CONFIG["model"],
            temperature=0.7
        )

        # 创建标题行识别的提示模板
        self.header_detection_template = PromptTemplate(
            input_variables=["description", "table_rows"],
            template="""请帮我分析下面这个表格的结构。标题行是指用于描述下方数据列含义的行，通常不是数据本身。

表格描述：
{description}

表格数据：
{table_rows}

请遵循以下步骤：
1. 仔细观察表格的结构。
2. 判断从第1行开始，连续有几行是作为整个表格的标题行（表头）。
3. 回答格式为：标题行数量：X行
"""
        )

        # 创建标题行识别LLM链
        self.header_detection_chain = LLMChain(
            llm=self.llm,
            prompt=self.header_detection_template
        )

        # 创建统计问答生成的提示模板
        self.statistical_prompt_template = PromptTemplate(
            input_variables=["table_header", "table_data"],
            template="""你是一个专业的数据分析专家。请分析以下完整表格，生成有意义的统计问答对。

表格标题：
{table_header}

完整表格数据：
{table_data}

要求：
1. 分析表格中的数值数据，生成包括多个数据值的统计类问答对（如最高值、最低值、平均值、合计、总数等）
2. 只生成结合多个单元格数据的统计值，不对单一情况下的具体单元格提问
3. 如果表格没有可统计的数值数据，可以不生成问答对
4. 问答对必须基于提供的数据，不能添加或修改任何信息
5. 主语不能使用"这""那"等代词，“表格”等代词需要说明表名
6. 统计问题应该具有实际意义和价值

请以标准JSON格式输出问答对：

{{
  "qa_pairs": [
    {{
      "question": " ",
      "answer": " "
    }},
    {{
      "question": " ",
      "answer": " "
    }}
  ]
}}

注意：
1. 只输出JSON格式的结果
2. 确保JSON格式正确，使用双引号
3. 问题和答案内容中的引号需要转义
4. 不要输出其他格式的内容
5. 如果没有有意义的统计值，返回空的qa_pairs数组
"""
        )

        # 创建统计问答生成LLM链
        self.statistical_qa_chain = LLMChain(
            llm=self.llm,
            prompt=self.statistical_prompt_template,
            output_parser=StatisticalQAOutputParser()
        )

    def detect_header_rows(self, table_rows: List[str], description: str = "") -> int:
        """使用大模型检测标题行数量（从第二行开始的标题行数量）"""
        try:
            # 取前5行进行分析，这些行已经是从第二行开始的
            analysis_rows = table_rows[:5]
            table_text = "\n".join([f"第{i+1}行: {row}" for i, row in enumerate(analysis_rows)])
            
            # 调用大模型分析
            result = self.header_detection_chain.run(
                description=description or "无描述",
                table_rows=table_text
            )
            # print(result)
            
            # 优先从"标题行数量：X行"格式中提取数字
            pattern_match = re.search(r'标题行数量[：:：]?\s*(\d+)\s*行?', result)
            if pattern_match:
                header_count = int(pattern_match.group(1))
            else:
                # 如果没有找到标准格式，尝试提取最后出现的数字
                numbers = re.findall(r'\d+', result)
                if numbers:
                    header_count = int(numbers[-1])  # 取最后一个数字
                else:
                    header_count = 1
            
            # 确保结果在合理范围内，最小为0（可能没有标题行），最大不超过可用行数
            header_count = max(0, min(header_count, len(table_rows)))
            
            print(f"    检测到标题行数量: {header_count}")
            return header_count
            
        except Exception as e:
            print(f"标题行检测失败，使用默认值1: {e}")
            return 1

    def extract_table_metadata(self, table_rows: List[str]) -> tuple:
        """提取表格的元数据：表名、描述、关键词、使用说明"""
        table_name = ""
        description = ""
        keywords = ""
        usage_instruction = ""
        
        try:
            if len(table_rows) >= 1:
                # 第一行包含所有元数据：表名、描述、关键词、使用说明
                first_cells = [cell.strip() for cell in table_rows[0].split(',')]
                
                # 第一列：表名
                if len(first_cells) >= 1 and first_cells[0]:
                    table_name = first_cells[0]
                
                # 第二列：描述
                if len(first_cells) >= 2 and first_cells[1]:
                    description = first_cells[1]
                
                # 第三列：关键词
                if len(first_cells) >= 3 and first_cells[2]:
                    keywords = first_cells[2]
                
                # 第四列：使用说明（可选）
                if len(first_cells) >= 4 and first_cells[3]:
                    usage_instruction = first_cells[3]
                
                # 调试输出，确认提取的内容
                print(f"    元数据提取调试:")
                print(f"      第一行原始内容: {table_rows[0][:100]}...")
                print(f"      分割后的单元格数量: {len(first_cells)}")
                print(f"      表名 (第1列): {table_name}")
                print(f"      描述 (第2列): {description[:50]}..." if description else "      描述 (第2列): 空")
                print(f"      关键词 (第3列): {keywords[:50]}..." if keywords else "      关键词 (第3列): 空")
                    
        except Exception as e:
            print(f"提取表格元数据时出错: {e}")
        
        return table_name, description, keywords, usage_instruction

    def generate_statistical_qa_from_table(self, table_data: Dict) -> List[Dict[str, str]]:
        """根据表格数据生成统计问答对"""
        try:
            # 获取表格元数据
            description = table_data.get('description', '')
            header = table_data['header']
            data_rows = table_data['data_rows']
            
            # 将数据行转换为字符串
            data_text = "\n".join(data_rows)
            
            # 创建增强的统计问答提示模板（只包含描述）
            enhanced_template = PromptTemplate(
                input_variables=["description", "header", "data_text"],
                template="""你是一个专业的统计分析专家。请根据以下表格数据，生成中文统计问答对。

表格描述：
{description}

表格标题行：
{header}

表格数据：
{data_text}

要求：
1. 分析表格中的数值数据，生成统计类问答对
2. 包括但不限于：总数、平均值、最大值、最小值、分布情况、趋势分析等
3. 结合表格描述来生成更有针对性的统计问答对
4. 问答对必须基于提供的数据，不能添加或修改任何信息
5. 主语不能使用"这""那"等代词
6. 确保统计结果的准确性

请以标准JSON格式输出问答对：

{{
  "qa_pairs": [
    {{
      "question": " ",
      "answer": " "
    }},
    {{
      "question": " ",
      "answer": " "
    }}
  ]
}}

注意：
1. 只输出JSON格式的结果
2. 确保JSON格式正确，使用双引号
3. 问题和答案内容中的引号需要转义
4. 不要输出其他格式的内容
"""
            )
            
            # 创建增强的LLM链
            enhanced_chain = LLMChain(
                llm=self.llm,
                prompt=enhanced_template,
                output_parser=StatisticalQAOutputParser()
            )
            
            result = enhanced_chain.run(
                description=description or "无描述",
                header=header,
                data_text=data_text
            )
            
            return result if isinstance(result, list) else []
        except Exception as e:
            print(f"生成统计问答对时出错: {e}")
            return []

    def parse_csv_tables(self, csv_file_path: str) -> List[Dict]:
        """解析CSV文件中以空行分隔的表格"""
        tables = []
        
        try:
            # 尝试多种编码格式读取CSV文件
            encodings = ['utf-8', 'gbk', 'gb2312', 'utf-8-sig', 'latin1']
            lines = None
            
            for encoding in encodings:
                try:
                    with open(csv_file_path, 'r', encoding=encoding) as f:
                        lines = f.readlines()
                    print(f"成功使用 {encoding} 编码读取文件")
                    break
                except UnicodeDecodeError:
                    continue
            
            if lines is None:
                print(f"无法使用任何编码格式读取文件: {csv_file_path}")
                return tables
            
            current_table = []
            
            for line in lines:
                line = line.strip()
                
                # 检查是否为空行或只包含分隔符的行
                is_empty_line = not line or re.match(r'^[,\s]*$', line)
                
                if is_empty_line:
                    if current_table:
                        # 处理当前表格
                        if len(current_table) >= 1:  # 至少需要包含元数据的第一行
                            # 提取表格元数据
                            table_name, description, keywords, usage_instruction = self.extract_table_metadata(current_table)
                            
                            print(f"    检测到表名: {table_name}")
                            print(f"    表格描述: {description[:50]}...")
                            print(f"    关键词: {keywords[:50]}...")
                            if usage_instruction:
                                print(f"    使用说明: {usage_instruction[:50]}...")
                            
                            # 从第2行开始检测标题行（跳过元数据行）
                            if len(current_table) >= 2:
                                header_detection_rows = current_table[1:6]  # 第2-6行用于标题行检测
                                if header_detection_rows:
                                    header_rows_count = self.detect_header_rows(header_detection_rows, description)
                                    # 标题行从第2行开始计算
                                    header_rows = current_table[1:1+header_rows_count]  # 第2行开始的标题行
                                    data_rows = current_table[1+header_rows_count:]     # 数据行
                                else:
                                    header_rows_count = 0
                                    header_rows = []
                                    data_rows = current_table[1:]  # 跳过第一行元数据
                            else:
                                header_rows_count = 0
                                header_rows = []
                                data_rows = []
                            
                            # 合并多行标题为一个标题字符串
                            header = "\n".join(header_rows)
                            
                            # 生成完整表格数据字符串（用于储存）
                            full_table_data = "\n".join(current_table)
                            
                            if data_rows:  # 确保有数据行
                                tables.append({
                                    'table_name': table_name,
                                    'description': description,
                                    'keywords': keywords,
                                    'usage_instruction': usage_instruction,
                                    'header': header,
                                    'data_rows': data_rows,
                                    'header_rows_count': header_rows_count,
                                    'full_table_data': full_table_data
                                })
                        current_table = []
                else:
                    # 过滤掉row_count和col_count信息
                    filtered_line = re.sub(r'\brow_count:\s*\d+', '', line)
                    filtered_line = re.sub(r'\bcol_count:\s*\d+', '', filtered_line)
                    # 清理多余的空格和逗号
                    filtered_line = re.sub(r'\s*,\s*,', ',', filtered_line)
                    filtered_line = re.sub(r'^\s*,\s*', '', filtered_line)
                    filtered_line = re.sub(r'\s*,\s*$', '', filtered_line)
                    filtered_line = filtered_line.strip()
                    
                    if filtered_line:
                        current_table.append(filtered_line)
            
            # 处理最后一个表格
            if current_table and len(current_table) >= 1:  # 至少需要包含元数据的第一行
                # 提取表格元数据
                table_name, description, keywords, usage_instruction = self.extract_table_metadata(current_table)
                
                print(f"    检测到表名: {table_name}")
                print(f"    表格描述: {description[:50]}...")
                print(f"    关键词: {keywords[:50]}...")
                if usage_instruction:
                    print(f"    使用说明: {usage_instruction[:50]}...")
                
                # 从第2行开始检测标题行（跳过元数据行）
                if len(current_table) >= 2:
                    header_detection_rows = current_table[1:6]  # 第2-6行用于标题行检测
                    if header_detection_rows:
                        header_rows_count = self.detect_header_rows(header_detection_rows, description)
                        # 标题行从第2行开始计算
                        header_rows = current_table[1:1+header_rows_count]  # 第2行开始的标题行
                        data_rows = current_table[1+header_rows_count:]     # 数据行
                    else:
                        header_rows_count = 0
                        header_rows = []
                        data_rows = current_table[1:]  # 跳过第一行元数据
                else:
                    header_rows_count = 0
                    header_rows = []
                    data_rows = []
                
                # 合并多行标题为一个标题字符串
                header = "\n".join(header_rows)
                
                # 生成完整表格数据字符串（用于储存）
                full_table_data = "\n".join(current_table)
                
                if data_rows:
                    tables.append({
                        'table_name': table_name,
                        'description': description,
                        'keywords': keywords,
                        'usage_instruction': usage_instruction,
                        'header': header,
                        'data_rows': data_rows,
                        'header_rows_count': header_rows_count,
                        'full_table_data': full_table_data
                    })
                
        except Exception as e:
            print(f"解析CSV文件时出错: {e}")
            
        return tables

    def process_csv_file(self, csv_file_path: str, output_file: str):
        """处理单个CSV文件，生成统计问答对"""
        print(f"正在处理文件: {os.path.basename(csv_file_path)}")
        
        # 解析CSV文件中的表格
        tables = self.parse_csv_tables(csv_file_path)
        
        if not tables:
            print(f"文件 {os.path.basename(csv_file_path)} 中未找到有效表格")
            return
        
        print(f"找到 {len(tables)} 个表格")
        
        all_qa_pairs = []
        
        for table_index, table in enumerate(tables, 1):
            table_name = table.get('table_name', '')
            description = table.get('description', '')
            keywords = table.get('keywords', '')
            full_table_data = table.get('full_table_data', '')
            
            table_info = f"表格 {table_index}"
            if table_name:
                table_info += f" (表名: {table_name})"
            table_info += f": 标题行数={table.get('header_rows_count', 0)}, 数据行数={len(table['data_rows'])}"
            print(f"  处理{table_info}")
            
            # 生成统计问答对
            qa_pairs = self.generate_statistical_qa_from_table(table)
            
            if qa_pairs:
                # 为每个问答对添加表格信息
                for qa_index, qa in enumerate(qa_pairs, 1):
                    qa['table_index'] = table_index
                    qa['qa_index_in_table'] = qa_index
                    qa['table_name'] = table_name
                    qa['table_description'] = description
                    qa['table_keywords'] = keywords
                    qa['table_header'] = table['header']
                    qa['qa_type'] = 'statistical'  # 标记为统计类问答对
                    
                    # 添加完整表格数据到context字段
                    context = {
                        "table_name": table_name,
                        "description": description,
                        "keywords": keywords,
                        "table_header": table['header'],
                        "data_rows": table['data_rows'],
                        "full_table_data": full_table_data,
                        "header_rows_count": table['header_rows_count']
                    }
                    qa['context'] = json.dumps(context, ensure_ascii=False)
                
                all_qa_pairs.extend(qa_pairs)
                print(f"    成功生成 {len(qa_pairs)} 个统计问答对")
            else:
                print(f"    表格 {table_index} 未生成统计问答对（可能没有可统计的数值数据）")
        
        # 保存到Excel文件
        if all_qa_pairs:
            category_id = os.path.splitext(os.path.basename(csv_file_path))[0]
            self.save_to_excel(all_qa_pairs, output_file, category_id)
            print(f"文件 {os.path.basename(csv_file_path)} 处理完成，共生成 {len(all_qa_pairs)} 个统计问答对")
        else:
            print(f"文件 {os.path.basename(csv_file_path)} 未生成任何统计问答对")

    def save_to_excel(self, qa_pairs: List[Dict[str, str]], filename: str, category_id: str = None):
        """保存问答对到Excel文件"""
        if not qa_pairs:
            print("没有问答对可保存")
            return
        
        # 准备数据
        data = []
        for i, qa in enumerate(qa_pairs, 1):
            row = {
                '标号': i,
                '提问': qa.get('question', ''),
                '回答': qa.get('answer', ''),
                '表名': qa.get('table_name', ''),
                '表格描述': qa.get('table_description', ''),
                '关键词': qa.get('table_keywords', ''),
                '表格标题': qa.get('table_header', ''),
                '问答类型': qa.get('qa_type', 'statistical')
            }
            if category_id:
                row['category_id'] = category_id
            if 'table_index' in qa:
                row['表格序号'] = qa['table_index']
            # 添加完整表格数据和context字段
            if 'context' in qa:
                row['context'] = qa['context']
            data.append(row)
        
        # 创建DataFrame
        columns = ['标号', '提问', '回答', '表名', '表格描述', '关键词', '表格标题', '问答类型']
        if category_id:
            columns.append('category_id')
        columns.extend(['表格序号', 'context'])
        
        df = pd.DataFrame(data, columns=columns)
        df.to_excel(filename, index=False)


def main():
    """主函数 - 批量处理CSV文件生成统计问答对"""
    generator = CSVStatisticalQAGenerator()
    
    print("=== CSV统计问答对生成工具 ===") 
    print(f"连接到阿里云百炼服务: {DASHSCOPE_CONFIG['base_url']}")
    print(f"使用模型: {DASHSCOPE_CONFIG['model']}")
    print()
    
    # 设置输入目录和输出目录（修改输出目录为独立的csvQAtotal文件夹）
    current_dir = os.path.dirname(os.path.abspath(__file__))
    input_dir = os.path.join(os.path.dirname(current_dir), "A_segmenter", "csv")
    output_dir = os.path.join(current_dir, "csvQAtotal")
    
    print(f"当前工作目录: {current_dir}")
    print(f"输入目录: {input_dir}")
    print(f"输出目录: {output_dir}")
    print()
    
    # 检查输入目录是否存在
    if not os.path.exists(input_dir):
        print(f"错误：输入目录不存在 - {input_dir}")
        return
    
    # 确保输出目录存在
    os.makedirs(output_dir, exist_ok=True)
    
    try:
        # 批量处理目录中的所有CSV文件
        generator.process_csv_directory(input_dir, output_dir)
        
    except Exception as e:
        print(f"批量处理过程中发生错误: {e}")


if __name__ == "__main__":
    main()